package com.zhongfu.ai.web;

import cn.hutool.core.util.IdUtil;
import com.zhongfu.ai.common.response.Result;
import org.redisson.api.RList;
import org.redisson.api.RedissonClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

@RestController
@RequestMapping("/api/rag")
@CrossOrigin("*")
public class RagController {

    private static final Logger log = LoggerFactory.getLogger(RagController.class);

    @Autowired
    private TokenTextSplitter tokenTextSplitter;

    @Autowired
    private PgVectorStore pgVectorStore;

    @Autowired
    private RedissonClient redissonClient;

    @GetMapping("/tag/list")
    public RList<String> tagList(){
        RList<String> tagList = redissonClient.getList("ragTag");
        return tagList;
    }

    @RequestMapping(value = "/file/upload", method = RequestMethod.POST)
    public Result upload(@RequestParam String ragTag,
                         @RequestParam("file") List<MultipartFile> files) {

        log.info("开始上传知识库: {}", ragTag);
        for(MultipartFile file : files) {
            TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(file.getResource());
            List<Document> documentList = tikaDocumentReader.get();

            List<Document> documentSpilterList = tokenTextSplitter.apply(documentList);
            // 添加知识库标签
            documentList.forEach(e->e.getMetadata().put("knowledge", ragTag));
            documentSpilterList.forEach(e->e.getMetadata().put("knowledge", ragTag));

            String fileUid = IdUtil.getSnowflakeNextIdStr();
            documentSpilterList.forEach(e->e.getMetadata().put("f_uid", fileUid));
            documentSpilterList.forEach(e->e.getMetadata().put("f_uid", fileUid));

            pgVectorStore.accept(documentSpilterList);

            // 添加知识库记录
            RList<String> tagList = redissonClient.getList("ragTag");
            if(!tagList.contains(ragTag)) {
                tagList.add(ragTag);
            }

            log.info("知识库上传完成：{}", ragTag);
        }

       return Result.success("1");
    }

}
